import torch
import torch.nn as nn
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

class MnistClassifer(nn.Module):
    def __init__(self):
        super().__init__()

        #define nerual network
        self.model = nn.Sequential(
            nn.Linear(784,300),
            #nn.ReLU(),
            #nn.Sigmoid(),
            nn.LeakyReLU(0.02),

            nn.Linear(300,100),
            #nn.ReLU(),
            #nn.Sigmoid(),
            nn.LeakyReLU(0.02),

            nn.Linear(100,10),
            #nn.ReLU(),
            #nn.Sigmoid(),
            nn.LeakyReLU(0.02) 
        )

        #create cost function
        self.loss_function = nn.MSELoss()
        #self.loss_function = nn.BCELoss()

        #create optimiser -- using SGD
        self.optimiser = torch.optim.SGD(self.parameters(),lr=0.02)

        #define training progress data
        self.counter = 0
        self.progress = []

    def forward(self,inputs):
        return self.model(inputs)
        
    def train(self,inputs,targets):
        #calculate dnn outputs
        outputs = self.forward(inputs)

        #calculate costs
        loss = self.loss_function(outputs,targets)

        #zero grad, backward propagation, update weights
        self.optimiser.zero_grad()
        loss.backward()
        self.optimiser.step()

        #show training progress
        self.counter += 1
        if (self.counter % 10 == 0):
            self.progress.append(loss.item())

            if (self.counter % 10000 == 0):
                print("counter = {0}, loss = {1}".format(self.counter,loss))

    def plot_progress(self):
        df = pd.DataFrame(self.progress, columns=['loss'])
        df.plot(ylim=(0,1.0), figsize=(16,8), alpha=0.1, marker=',', grid=True, yticks=(0,0.25,0.5))
        plt.show()

def train():      
    #download data
    df = pd.read_csv('./mnist_data/mnist_train.csv',header=None)

    #define network
    net = MnistClassifer()
    #training 10 epochs
    epochs = 10

    for i in range(epochs):
        print("training epoch {0} of {1}".format(i+1, epochs))

        num_rows = len(df)
        for i in range(num_rows):
            data = df.iloc[i]

            #acquire label
            label = data[0]
            target_tensor = torch.zeros((10))
            target_tensor[label] = 1.0

            image_data_tensor = torch.FloatTensor(data[1:].values)/255.0

            net.train(image_data_tensor,target_tensor)

    torch.save(net,'mnist_parameters.pt')
    net.plot_progress()


def test():
    # load net parameters
    net = torch.load('mnist_parameters.pt')
    # load test data
    df = pd.read_csv('./mnist_data/mnist_test.csv',header=None)
    err = 0
    
    for i in range(1000):
        data = df.iloc[i]
        label = data[0]
        image_data_tensor = torch.FloatTensor(data[1:].values)/255.0
        # test 
        output = net(image_data_tensor)
        output_np = output.detach().numpy()
        output_label = np.argmax(output_np)
        print(output_np)
        print(output_label)
        if output_label != label:
            err += 1
    print(f'error={err}')
    

test()